"""
训练模型
"""

from parl.algorithms import DDQN
from parl.utils import ReplayMemory

from Examples.DQNExamples.DimensionOne.DimOneAgent import DimOneAgent
from Examples.DQNExamples.DimensionOne.DimOneModel import DimOneModel
from Examples.DQNExamples.SlayTheSpire.SlayTheSpireBase import Player
from Examples.DQNExamples.SlayTheSpire.SlayTheSpireCards import all_cards
from Examples.DQNExamples.SlayTheSpire.SlayTheSpireEnv import SlayTheSpireEnv
from Examples.DQNExamples.SlayTheSpire.SlayTheSpireMonster import Cultist
from Examples.DQNExamples.SlayTheSpire.SlayTheSpireTraining import SlayTheSpireTraining

MEMORY_SIZE = 200000  # 记忆库大小
BATCH_SIZE = 32  # 批大小
GAMMA = 0.999  # 奖励衰减率
LEARNING_RATE = 0.0001  # 学习率
E_GREED = 0.8  # 贪心率
E_GREED_DECREMENT = 1e-9  # 贪心率衰减率
UPDATE_TARGET_STEPS = 200  # 更新目标网络的步数
MAX_EPISODES = 10000000  # 最大回合数
MAX_STEPS = 4000  # 训练每回合最大步数, 避免训练时间过长
MAX_TEST_STEPS = 2000  # 测试每回合最大步数, 避免时间过长
MEMORY_WARMUP_SIZE = 4000  # replay memory预热大小, 需要预存一些经验数据后才开始训练
TEST_FREQ = 2000  # 测试的频率, 每隔多少回合测试一次
FRAME_SKIPPING = 1  # 每隔多少帧采样一次
RENDER = False  # 训练是否渲染
TEST_RENDER = True  # 测试是否渲染
TARGET_UPDATE_FREQ = 5  # 每隔多少次训练更新一次target_net
SAVE_PATH = 'model/slayTheSpire.ckpt'  # 模型保存路径

# 创建环境
player = Player(max_hp=100, max_mp=3)
monster = Cultist()
cards = all_cards
env = SlayTheSpireEnv(player=player, monsters=[monster], cards=cards)

# 获取动作维度和状态维度
action_dim = 11
# obs_dim = env.observation_space.n
obs_dim = 255

# 经验回放池
rpm = ReplayMemory(MEMORY_SIZE, obs_dim, 0)

# 创建模型
model = DimOneModel(act_dim=action_dim, obs_dim=obs_dim)

# 创建算法
algorithm = DDQN(model, gamma=GAMMA, lr=LEARNING_RATE)

# 创建Agent
agent = DimOneAgent(algorithm, obs_dim, action_dim,
                    e_greed=E_GREED,
                    e_greed_decrement=E_GREED_DECREMENT,
                    update_target_steps=UPDATE_TARGET_STEPS)

# 创建训练
train = SlayTheSpireTraining()

# 开始训练
train.train(
    agent=agent,
    env=env,
    rpm=rpm,
    batch_size=BATCH_SIZE,
    max_steps=MAX_STEPS,
    max_test_steps=MAX_TEST_STEPS,
    memory_warmup_size=MEMORY_WARMUP_SIZE,
    test_freq=TEST_FREQ,
    render=RENDER,
    test_render=TEST_RENDER,
    max_episodes=MAX_EPISODES,
    frame_skipping=FRAME_SKIPPING,
    save_path=SAVE_PATH,
)
